package com.uestc.utils;

import com.uestc.entity.HuffmanNode;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;

public class HuffmanCoding {
    public static HuffmanNode buildHuffmanTree(Map<Byte, Integer> freqMap) {
        PriorityQueue<HuffmanNode> pq = new PriorityQueue<>();
        for (Map.Entry<Byte, Integer> entry : freqMap.entrySet()) {
            pq.add(new HuffmanNode(entry.getKey(), entry.getValue()));
        }

        while (pq.size() != 1) {
            HuffmanNode left = pq.poll();
            HuffmanNode right = pq.poll();

            HuffmanNode sum = new HuffmanNode(Byte.valueOf((byte) 10), left.freq + right.freq);
            sum.isLeaf = false;
            sum.left = left;
            sum.right = right;

            pq.add(sum);
        }

        return pq.peek();
    }

    public static void encode(HuffmanNode root, String str, Map<Byte, String> huffmanCode) {
        if (root == null) {
            return;
        }

        if (root.isLeaf) {
            huffmanCode.put(root.data, str);
        }

        encode(root.left, str + "0", huffmanCode);
        encode(root.right, str + "1", huffmanCode);
    }

    public static String encodeData(byte[] data, Map<Byte, String> huffmanCode) {
        StringBuilder sb = new StringBuilder();
        for (Byte c : data) {
            sb.append(huffmanCode.get(c));
        }
        return sb.toString();
    }

    public static List<Byte> decodeData(HuffmanNode root, String encodedData) {
        List<Byte> bytes = new ArrayList<>();
        HuffmanNode curr = root;
        for (char bit : encodedData.toCharArray()) {
            if (bit == '0') {
                curr = curr.left;
            } else {
                curr = curr.right;
            }

            if (curr.left == null && curr.right == null) {
                bytes.add(curr.data);
                curr = root;
            }
        }
        return bytes;
    }

    public static Map<Byte, Integer> buildFrequencyMap(byte[] data) {
        Map<Byte, Integer> freqMap = new HashMap<>();
        for (byte c : data) {
            freqMap.put(c, freqMap.getOrDefault(c, 0) + 1);
        }
        return freqMap;
    }

    public static void compress(List<File> sourceFiles, String outputPath) {

        try (DataOutputStream dos = new DataOutputStream(new FileOutputStream(outputPath))) {
            dos.writeInt(sourceFiles.size());
            for (File file : sourceFiles) {
                byte[] bytes = Files.readAllBytes(Paths.get(file.getPath()));
                Map<Byte, Integer> freqMap = buildFrequencyMap(bytes);
                HuffmanNode root = buildHuffmanTree(freqMap);
                Map<Byte, String> huffmanCode = new HashMap<>();
                encode(root, "", huffmanCode);
                String encodedData = encodeData(bytes, huffmanCode);
                writeOneFile(file.getName(), dos, freqMap, encodedData);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private static void writeOneFile(String fileName, DataOutputStream dos, Map<Byte, Integer> freqMap, String encodedData) throws IOException {
        //写入文件名
        dos.writeInt(fileName.length());
        dos.writeChars(fileName);

        dos.writeInt(freqMap.size());
        for (Map.Entry<Byte, Integer> entry : freqMap.entrySet()) {
            dos.writeByte(entry.getKey());
            dos.writeInt(entry.getValue());
        }

        dos.writeInt(encodedData.length());

        int p = 0;
        List<Byte> bytes = new ArrayList<>();
        while (p + 8 < encodedData.length()) {
            byte b = 0x00;
            for (int i = 0; i < 8; p++, i++) {
                char c = encodedData.charAt(p);
                if (c == '1') {
                    b = (byte) ((b << 1) + 1);
                } else {
                    b = (byte) (b << 1);
                }
            }
            bytes.add(b);
        }

        byte b = 0x00;
        for (int i = p; i < encodedData.length(); i++) {
            char c = encodedData.charAt(i);
            if (c == '1') {
                b = (byte) ((b << 1) + 1);
            } else {
                b = (byte) (b << 1);
            }
        }
        int diff = 8 + p - encodedData.length();
        b = (byte) (b << diff);

        bytes.add(b);

        byte[] bytesArray = new byte[bytes.size()];
        for (int i = 0; i < bytes.size(); i++) {
            bytesArray[i] = bytes.get(i);
        }
        dos.write(bytesArray);
    }

    public static void decompress(List<File> sourceFiles, String outputPath) {
        for (File file : sourceFiles) {
            try (DataInputStream dis = new DataInputStream(new FileInputStream(file.getPath()))) {
                int fileNum = dis.readInt();
                for (int i = 0; i < fileNum; i++) {
                    readOneFile(outputPath, dis);
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public static void decompress(byte[] bytes, String outputPath) {
        try (DataInputStream dis = new DataInputStream(new ByteArrayInputStream(bytes))) {
            int fileNum = dis.readInt();
            for (int i = 0; i < fileNum; i++) {
                readOneFile(outputPath, dis);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    private static void readOneFile(String outputPath, DataInputStream dis) throws IOException {
        int fileNameLength = dis.readInt();
        StringBuilder fileName = new StringBuilder();
        for (int i = 0; i < fileNameLength; i++) {
            fileName.append(dis.readChar());
        }

        int freqMapSize = dis.readInt();
        Map<Byte, Integer> freqMap = new HashMap<>();
        for (int i = 0; i < freqMapSize; i++) {
            Byte key = dis.readByte();
            int value = dis.readInt();
            freqMap.put(key, value);
        }

        int encodedDataLength = dis.readInt();
        double divisionResult = (double) encodedDataLength / 8;
        int bytesCount = (int) Math.ceil(divisionResult);
        byte[] fileBytes = new byte[bytesCount];
        dis.readFully(fileBytes);
        StringBuilder encodedData = new StringBuilder();
        for (int i = 0; i < bytesCount; i++) {
            encodedData.append(ByteToBinaryString.byteToBinaryString(fileBytes[i]));
        }

        HuffmanNode root = buildHuffmanTree(freqMap);
        List<Byte> decodedData = decodeData(root, encodedData.substring(0, encodedDataLength));
        byte[] bytes = new byte[decodedData.size()];
        for (int i = 0; i < decodedData.size(); i++) {
            bytes[i] = decodedData.get(i);
        }

        Files.write(Paths.get(outputPath + "\\" + fileName), bytes);
    }

    public static void main(String[] args) {
        String inputPath = "test.txt";
        List<File> files = new ArrayList<>();
        files.add(new File(inputPath));
        String compressedPath = "compressed.huff";
        String decompressedPath = "F:/桌面";

        compress(files, compressedPath);

        List<File> compressedFiles = new ArrayList<>();
        compressedFiles.add(new File(compressedPath));
        decompress(compressedFiles, decompressedPath);

        System.out.println("Compression and decompression completed successfully.");
    }
}